#pragma once
#include "vector.h"
#include <vector>

// Stencil that computes the Laplacian at a point.
struct LaplacianStencil {
    int xp, xm, yp, ym, zp, zm;
};
struct IJK { 
    int i,j,k; 
}; 

class Solver {
public:
    // Constructor.
    Solver(int nx, int ny, int nz);
    // Computes a forward Euler timestep.
    void fwd_euler_step(int nsteps, double dt);
    void cn_step(int nsteps, double dt);

    void init();

    IJK to_ijk(int i) const;
    int from_ijk(const IJK &ijk) const;
    int size() const { return _T.size(); }

private:
    int _nx, _ny, _nz;
    Vector _T, _T_next; 
    double _alpha, _idx2, _idy2, _idz2;
    std::vector<LaplacianStencil> _stencil;
};

// Functions to go from a flat index to ijk index, and back.
inline IJK Solver::to_ijk(int i) const {
    return {i%_nx, (i/_nx)%_ny, i/_nx/_ny};
}
inline int Solver::from_ijk(const IJK &ijk) const {
    return ijk.i + ijk.j*_nx + ijk.k*_nx*_ny;
}

